
import time
import queue
from coptur.spider import Spider
from coptur.scanner import Scanner
from coptur.models.task import Task


class Worker:
    def __init__(self, options):
        self.options = options
        self.build_task() # => build self.task
        self.build_queue() # build self.site_queue & self.host_queue
        self.build_threads() # build self.worker_threads

    def run(self):
        if self.options.scan_type == 1: # host vuln only
            self.run_port_scan()
            #self.run_host_scan()
        elif self.options.scan_type == 2: # web vuln only
            self.run_web_spider()
            self.run_web_scan()
        elif self.options.scan_type == 3: # url vuln only
            # do not run spider
            self.run_web_scan()
        else: # both
            # 将来开三个进程来运行下面的任务
            self.run_host_scan()
            self.run_web_spider()
            self.run_web_scan()


    # 主机漏洞扫描
    def run_host_scan(self):
        self.host_scanners = []
        for thread in self.host_scan_threads:
            scanner = Scanner( thread, "host", self.host_queue )
            scanner.start()
            self.host_scanners.append(scanner)

        for s in self.host_scanners:
            s.join()



    def run_web_scan(self):
        pass


    # 对任务当中的主机进行端口开放扫描
    def run_port_scan(self):
        self.port_scanners = []
        for thread in self.port_scan_threads:
            scanner = Scanner( thread, "port", self.host_queue )
            scanner.start()
            self.port_scanners.append(scanner)

        for s in self.port_scanners:
            s.join()

    # 使用爬虫对目标进行递归爬行，爬行到的urls存入数据库
    def run_web_spider(self):
        self.spiders = []
        for thread in self.worker_threads:
            spider = Spider( thread, self.worker_queue )
            spider.start()
            self.spiders.append(spider)

        for s in self.spiders:
            s.join()



    # 创建任务数据对象
    def build_task(self):
        self.task = Task( scan_type=self.options.scan_type,
                          target = self.options.targets,
                          options=str(self.options)
                          )


    # 创建扫描队列
    def build_queue(self):
        #无阻塞队列
        self.host_queue = queue.Queue()
        self.site_queue = queue.Queue()

        #扫描目标以逗号分隔
        for target in self.options.targets.split(","):
            #print( "填充扫描对象 - %s" % target )
            self.host_queue.put(target)
            self.site_queue.put(target)




    # 创建扫描线程
    def build_threads(self):
        self.host_scan_threads = []
        self.port_scan_threads = []
        self.web_spider_threads = []
        self.web_scan_threads = []
        for i in range( 0,self.options.threads ):
            #print( "构建扫描线程 - %d" % i )
            self.host_scan_threads.append( "主机扫描线程-%d" % i )
            self.port_scan_threads.append( "端口扫描线程-%d" % i )
            self.web_spider_threads.append( "Web爬虫线程-%d" % i )
            self.web_scan_threads.append( "Web扫描线程-%d" % i )
